fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed install by ganeshkumarashok · Pull Request #8919 · Azure/AgentBaker

ganeshkumarashok · 2026-07-13T21:21:59Z

Summary

Fixes nvidia-smi failing with "Failed to initialize NVML: Driver/library version mismatch" on agentpool GRID / converged A10 (NVv5) nodes and NAP GRID nodes booting the CUDA-prebaked shared Ubuntu VHD.

Root cause: the shared Ubuntu VHD prebakes only the cuda(-lts) driver + a DKMS marker (#8786/#8803). The existing cleanUpPrebakedGPUDriver only runs on nodes that do not install a managed driver (GPU_NODE != true OR skip). A GRID node installs the grid driver on top of the stale prebaked cuda module + /usr/bin/lib64 libs, and the two collide → NVML mismatch.

Fix

Add cleanUpGridNodeCudaPrebake, invoked from ensureGPUDrivers before the configGPUDrivers / validateGPUDrivers dispatch (so both install paths are covered). It tears down the prebake when:

this node's driver type maps to grid (grid / grid-v20), and
the prebake marker is not grid — i.e. a cuda marker, or a legacy marker with no driver_kind= line.

Pure driver-kind mismatch, so no version comparison is needed. Reuses cleanUpPrebakedGPUDriver for removal (also clears /usr/bin/lib64 + the DKMS registration). No-op when there is no prebake marker (older / non-prebake VHDs) or on a CUDA node.

Scope

GRID mismatch only. CUDA nodes are intentionally untouched — the NAP cuda same-kind / different-version case is a separate issue, addressed by aligning the driver version NAP requests with the baked image.

Branch strategy

This is the main-targeted PR (source of truth) so future VHDs cut from main carry the fix. The same commit is cherry-picked to the live prebake-enabled release branch(es) as separate backport PRs (starting with official/v20260702, #8918).

Test plan

shellcheck clean on the new function
Isolated logic harness: teardown on grid+cuda-marker, grid-v20+cuda, grid+legacy-marker; no-op on cuda-lts node, grid+grid-marker, no marker, non-Ubuntu
ShellSpec coverage in cse_config_spec.sh (cleanUpGridNodeCudaPrebake)
e2e on an agentpool GRID/A10 node + a NAP GRID node booting the prebaked VHD

🤖 Generated with Claude Code

… install The shared Ubuntu VHD prebakes only the cuda(-lts) driver plus a DKMS marker. A GRID / converged (A10, NVv5) node then installs the grid driver on top, and the stale prebaked cuda module + its /usr/bin/lib64 userspace libs collide with the grid driver, so nvidia-smi fails with "Failed to initialize NVML: Driver/library version mismatch". The existing cleanUpPrebakedGPUDriver only runs on nodes that do NOT install a managed driver (GPU_NODE != true OR skip), so managed GRID nodes keep the stale prebake. This adds cleanUpGridNodeCudaPrebake, run from ensureGPUDrivers before the configGPUDrivers/validateGPUDrivers dispatch (so both paths are covered), which tears down the prebake when this node installs grid and the marker is not grid (a cuda or legacy kind-less marker). Pure driver-kind mismatch, so no version comparison is needed. CUDA nodes are intentionally untouched. Scoped to the GRID mismatch only; the NAP cuda same-kind/different-version case is addressed separately by aligning the driver version NAP requests with the baked image.

Copilot

Pull request overview

Adds a targeted cleanup step for Ubuntu GRID nodes to remove a mismatched CUDA(-lts) driver prebake marker/install before the managed GRID driver path runs, preventing NVML “Driver/library version mismatch” failures on GRID / NVv5 (A10) nodes booting the shared prebaked Ubuntu VHD.

Changes:

Introduces cleanUpGridNodeCudaPrebake and invokes it early in ensureGPUDrivers on Ubuntu before the config/validate dispatch.
Uses the existing prebake DKMS marker to decide teardown (only keep when marker kind is explicitly grid).
Adds ShellSpec coverage for the teardown decision logic.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
parts/linux/cloud-init/artifacts/cse_config.sh	Adds GRID-vs-CUDA prebake mismatch teardown logic and hooks it into `ensureGPUDrivers` for Ubuntu.
spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh	Adds ShellSpec tests validating teardown/no-op decisions for `cleanUpGridNodeCudaPrebake`.

ganeshkumarashok · 2026-07-13T21:44:05Z

+        It 'is a no-op when no prebake marker exists'
+            GPU_DKMS_MARKER_FILE="$(mktemp)"; rm -f "${GPU_DKMS_MARKER_FILE}"
+            NVIDIA_GPU_DRIVER_TYPE="grid"
+            When call cleanUpGridNodeCudaPrebake
+            The output should not include "STUB_TEARDOWN_CALLED"
+            The status should be success
+        End
+    End


Good catch — added a non-Ubuntu no-op test in 57123d6. It sets OS=MARINER with a mismatched cuda marker + grid node and asserts cleanUpPrebakedGPUDriver is not called (and exit 0), guarding the OS != UBUNTU_OS_NAME early return against accidental removal/inversion. Mirrored to the 0702 backport (#8918).

ganeshkumarashok · 2026-07-13T21:36:35Z

Note on CI: the lint (aks-node-controller) check failure is pre-existing on main and unrelated to this PR.

Error: aks-node-controller/parser/helper.go:553:10: string "AzureChinaCloud" has 3 occurrences, make it a constant (goconst)
This PR touches only parts/linux/cloud-init/artifacts/cse_config.sh and its ShellSpec — zero aks-node-controller/ files (and no Go at all).
The flagged code already exists on origin/main at that line; the lint job runs with only-new-issues: false, so it reports the whole module's existing issues, not just this diff's.

All checks relevant to this change pass: shellcheck, shellspec, go-test, Pester, cue, dcgm-compatibility. Leaving the goconst fix out of scope for this hotfix.

Addresses Copilot review on #8919: cover the OS != UBUNTU_OS_NAME early return so a regression that removes or inverts the Ubuntu gate is caught.

Copilot

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

…to v20260702] (#8921) Co-authored-by: aks-node-assistant[bot] <190555641+aks-node-assistant[bot]@users.noreply.github.com>

…to v20260709] (#8924) Co-authored-by: aks-node-assistant[bot] <190555641+aks-node-assistant[bot]@users.noreply.github.com>

Copilot AI review requested due to automatic review settings July 13, 2026 21:22

ganeshkumarashok requested review from AbelHu, Devinwong, SriHarsha001, awesomenix, calvin197, cameronmeissner, djsly, karenychen, lilypan26, mxj220, pdamianov-dev, phealy, r2k1, runzhen, sulixu, surajssd, timmy-wright, titilambert, xuexu6666 and zachary-bailey as code owners July 13, 2026 21:22

Copilot started reviewing on behalf of ganeshkumarashok July 13, 2026 21:23 View session

sulixu approved these changes Jul 13, 2026

View reviewed changes

Copilot AI reviewed Jul 13, 2026

View reviewed changes

test(gpu): add non-Ubuntu no-op case for cleanUpGridNodeCudaPrebake

57123d6

Addresses Copilot review on #8919: cover the OS != UBUNTU_OS_NAME early return so a regression that removes or inverts the Ubuntu gate is caught.

Copilot AI review requested due to automatic review settings July 13, 2026 21:42

Copilot started reviewing on behalf of ganeshkumarashok July 13, 2026 21:46 View session

Copilot AI reviewed Jul 13, 2026

View reviewed changes

sulixu approved these changes Jul 13, 2026

View reviewed changes

ganeshkumarashok merged commit 947d2de into main Jul 13, 2026
20 of 21 checks passed

ganeshkumarashok deleted the ganesh/grid-prebake-teardown-main branch July 13, 2026 21:52

This was referenced Jul 13, 2026

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes [backport #8919 to v20260702] #8921

Merged

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed install #8918

Closed

ganeshkumarashok added a commit that referenced this pull request Jul 13, 2026

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes [backport #8919 …

bfd5a47

…to v20260702] (#8921) Co-authored-by: aks-node-assistant[bot] <190555641+aks-node-assistant[bot]@users.noreply.github.com>

ganeshkumarashok mentioned this pull request Jul 13, 2026

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes [backport #8919 to v20260709] #8924

Merged

2 tasks

awesomenix mentioned this pull request Jul 13, 2026

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed… #8925

Closed

ganeshkumarashok added a commit that referenced this pull request Jul 13, 2026

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes [backport #8919 …

da3471a

…to v20260709] (#8924) Co-authored-by: aks-node-assistant[bot] <190555641+aks-node-assistant[bot]@users.noreply.github.com>

djsly mentioned this pull request Jul 14, 2026

fix: clean up mismatched prebaked GPU drivers #8915

Closed

This was referenced Jul 15, 2026

fix(gpu): tear down prebaked GPU driver on version mismatch, not just kind mismatch #8945

Open

fix(gpu): unload full nvidia module stack in prebake cleanup, not just idle bare module #8955

Open

Navigation Menu

Uh oh!

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed install - #8919

fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed install#8919
ganeshkumarashok merged 2 commits into
mainfrom
ganesh/grid-prebake-teardown-main

ganeshkumarashok commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

ganeshkumarashok Jul 13, 2026

Uh oh!

ganeshkumarashok commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

ganeshkumarashok commented Jul 13, 2026

Summary

Fix

Scope

Branch strategy

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

ganeshkumarashok Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

ganeshkumarashok commented Jul 13, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants